{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.simplefilter(action='ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import umap\n",
    "from collections import OrderedDict\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tnrange, tqdm_notebook\n",
    "import deepchem as dc\n",
    "from scipy.spatial import distance\n",
    "from sklearn.metrics import roc_auc_score, average_precision_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change working directory\n",
    "os.chdir('/home/yuke/PythonProject/DrugEmbedding/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hvae import *\n",
    "from evae import *\n",
    "from drugdata import *\n",
    "from metrics import *\n",
    "from lorentz import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from decode import *\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "from rdkit import rdBase\n",
    "rdBase.DisableLog('rdApp.error') #disable RDKit warning messages\n",
    "from rdkit.Chem.Draw import IPythonConsole #Needed to show molecules\n",
    "from rdkit import DataStructs\n",
    "from rdkit.Chem.Fingerprints import FingerprintMols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def euc_dist(z1, z2):\n",
    "    return np.linalg.norm(z1-z2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def angular_dist(z1, z2):\n",
    "    cos_sim = np.dot(z1,z2)/(np.linalg.norm(z1) * np.linalg.norm(z2))\n",
    "    return np.arccos(cos_sim)/np.pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lor_dist(z1, z2):\n",
    "    m = z1*z2\n",
    "    lor_prod = m[1:].sum() - m[0]\n",
    "    x = - lor_prod\n",
    "    x = np.where(x<1.0, 1.0+1e-6, x)\n",
    "    return np.log(x + np.sqrt(x**2 - 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lor_pairwise_dist(x):\n",
    "    x0 = x[:,0].reshape(-1,1)\n",
    "    x1 = x[:,1:]\n",
    "    m = np.matmul(x1, x1.transpose()) - np.matmul(x0, x0.transpose())\n",
    "    np.fill_diagonal(m, -1-1e-12)\n",
    "    m = -m\n",
    "    m = np.where(m<1.0, 1.0 + 1e-6, m)\n",
    "    dm = np.log(m + np.sqrt(m**2 - 1))\n",
    "    # prevent Inf. distance in dm\n",
    "    dm = np.nan_to_num(dm)\n",
    "    return dm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_durg_atc(ATC_LVL):\n",
    "    drug_atc_dict = {}\n",
    "    for i in range(len(drug_lst)):\n",
    "        d = drug_lst[i]\n",
    "        drug_atc_dict[i] = {d:df_atc[df_atc['ATC_LVL5']==d][ATC_LVL].unique()}\n",
    "    return drug_atc_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_drug_repurpose(drug_x, ATC_LVL, dm, drug_lst):\n",
    "    # step 1: create drug atc dictionary\n",
    "    drug_atc_dict = get_durg_atc(ATC_LVL)\n",
    "    # step 2: get drug2drug distance array\n",
    "    drug_drug_dist = dm[drug_lst.index(drug_x),:]\n",
    "    # create atc level drug2drug distance dataframe\n",
    "    atc_drug_dist = []\n",
    "    for i in range(len(drug_lst)):\n",
    "        d = drug_lst[i]\n",
    "        if d != drug_x:\n",
    "            atc_array = drug_atc_dict[i][d]\n",
    "            if len(atc_array) > 0:\n",
    "                for atc in atc_array:\n",
    "                    atc_drug_dist.append((atc, drug_drug_dist[i]))\n",
    "    labels = [ATC_LVL, 'Drug2Drug_Dist']\n",
    "    df = pd.DataFrame.from_records(atc_drug_dist, columns=labels)\n",
    "    df_agg = df.groupby([ATC_LVL]).agg(['mean', 'std', 'count'])\n",
    "    df_agg.columns = ['_'.join(col).strip() for col in df_agg.columns.values]\n",
    "    return df_agg.reset_index().sort_values('Drug2Drug_Dist_mean')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def np_lorentz_product(z1, z2):\n",
    "    m = z1 * z2\n",
    "    return m[1:].sum() - m[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def np_arccoh(x):\n",
    "    x = np.where(x<1.0, 1.0 + 1e-6, x)\n",
    "    return np.log(x + np.sqrt(x**2 - 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def np_inv_exp_map(z, mu):\n",
    "    alpha = - np_lorentz_product(mu, z)\n",
    "    return (np_arccoh(alpha)/np.sqrt(alpha**2 - 1)) * (z - alpha*mu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_dir = './experiments/KDD/kdd_010'\n",
    "checkpoint = 'checkpoint_epoch110.model'\n",
    "config_path = os.path.join(exp_dir, 'configs.json')\n",
    "checkpoint_path = os.path.join(exp_dir, checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(config_path, 'r') as fp:\n",
    "    configs = json.load(fp)\n",
    "fp.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configs['checkpoint'] = checkpoint\n",
    "configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_model(configs)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load SMILES Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 625,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = OrderedDict()\n",
    "splits = ['train', 'valid', 'test']\n",
    "for split in splits:\n",
    "    datasets[split] = drugdata(task = configs['task'],\n",
    "                               fda_drugs_dir=configs['data_dir'],\n",
    "                               fda_smiles_file=configs['fda_file'],\n",
    "                               fda_vocab_file=configs['vocab_file'],\n",
    "                               fda_drugs_sp_file=configs['atc_sim_file'],\n",
    "                               experiment_dir=os.path.join(configs['checkpoint_dir'], configs['experiment_name']),\n",
    "                               smi_file='smiles_' + split + '.smi',\n",
    "                               max_sequence_length=configs['max_sequence_length'],\n",
    "                               nneg=configs['nneg'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 626,
   "metadata": {},
   "outputs": [],
   "source": [
    "fda_drug_lst = []\n",
    "fda_smiles_lst = []\n",
    "\n",
    "# retrieve SMILES from dataloader\n",
    "for split in ['train', 'valid', 'test']:\n",
    "    dd = datasets[split]\n",
    "    for k,v in dd.smiles.items():\n",
    "        if k[:4] != 'zinc':\n",
    "            smi = ''\n",
    "            token_lst = v['words'][1:] # skip <sos>\n",
    "            for t in token_lst:\n",
    "                if t != '<eos>':\n",
    "                    smi = smi + t \n",
    "                else:\n",
    "                    break\n",
    "            fda_drug_lst.append(k)\n",
    "            fda_smiles_lst.append(smi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 627,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [03:31<00:00, 21.18s/it]\n"
     ]
    }
   ],
   "source": [
    "drug_lst_train, mean_lst_train, logv_lst_train = fda_drug_rep(configs, datasets['train'], model, all_drugs=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 628,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:11<00:00, 11.24s/it]\n"
     ]
    }
   ],
   "source": [
    "drug_lst_valid, mean_lst_valid, logv_lst_valid = fda_drug_rep(configs, datasets['valid'], model, all_drugs=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 629,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:14<00:00, 14.74s/it]\n"
     ]
    }
   ],
   "source": [
    "drug_lst_test, mean_lst_test, logv_lst_test = fda_drug_rep(configs, datasets['test'], model, all_drugs=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 630,
   "metadata": {},
   "outputs": [],
   "source": [
    "drug_lst = drug_lst_train + drug_lst_valid + drug_lst_test\n",
    "assert drug_lst == fda_drug_lst, 'FDA drug names do not match!'\n",
    "mean_lst = mean_lst_train + mean_lst_valid + mean_lst_test\n",
    "logv_lst = logv_lst_train + logv_lst_valid + logv_lst_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 631,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert list to numpy array\n",
    "for i in range(len(drug_lst)):\n",
    "    mean_lst[i] = np.array(mean_lst[i])\n",
    "    logv_lst[i] = np.array(logv_lst[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 632,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create SMILES dataframe\n",
    "df_mdl = pd.DataFrame({'ATC_LVL5': drug_lst, 'smiles': fda_smiles_lst, 'mu': mean_lst, 'logv': logv_lst})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 633,
   "metadata": {},
   "outputs": [],
   "source": [
    "## remove drugs with duplicated entries\n",
    "#idx = df_mdl.groupby('ATC_LVL5').size() > 1\n",
    "#drug_lst_clean = list(df_mdl.groupby('ATC_LVL5').size()[-idx].reset_index()['ATC_LVL5'])\n",
    "#df_smiles = df_mdl[df_mdl['ATC_LVL5'].isin(drug_lst_clean)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 634,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load df_fda_year\n",
    "df_fda_year = pd.read_csv('./experiments/EXP_TASK/exp_task_011/all_drugs_approval_year.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 635,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_smiles = df_mdl.merge(df_fda_year, how='left', left_on='ATC_LVL5', right_on='drug_name')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 636,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ATC_LVL5</th>\n",
       "      <th>smiles</th>\n",
       "      <th>mu</th>\n",
       "      <th>logv</th>\n",
       "      <th>drug_name</th>\n",
       "      <th>Approval</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>novobiocin</td>\n",
       "      <td>CO[C@@H]1[C@@H](OC(N)=O)[C@@H](O)[C@H](Oc2ccc3...</td>\n",
       "      <td>[13.137248039245605, 1.3716269731521606, 0.025...</td>\n",
       "      <td>[-6.649742126464844, -5.984184741973877, -9.82...</td>\n",
       "      <td>novobiocin</td>\n",
       "      <td>1964</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>vildagliptin</td>\n",
       "      <td>N#C[C@@H]1CCCN1C(=O)CNC12CC3CC(CC(O)(C3)C1)C2</td>\n",
       "      <td>[102.30171203613281, 2.743731737136841, -1.416...</td>\n",
       "      <td>[-4.914798259735107, -4.87358283996582, -7.162...</td>\n",
       "      <td>vildagliptin</td>\n",
       "      <td>2007</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>benazepril</td>\n",
       "      <td>CCOC(=O)[C@H](CCc1ccccc1)N[C@H]1CCc2ccccc2N(CC...</td>\n",
       "      <td>[146.45880126953125, 1.0793601274490356, 2.399...</td>\n",
       "      <td>[-4.702251434326172, -4.654637813568115, -6.57...</td>\n",
       "      <td>benazepril</td>\n",
       "      <td>1991</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>testosterone</td>\n",
       "      <td>CCC(=O)O[C@H]1CC[C@H]2[C@@H]3CCC4=CC(=O)CC[C@]...</td>\n",
       "      <td>[121.25065612792969, 1.7730571031570435, 0.622...</td>\n",
       "      <td>[-4.358664035797119, -4.222125053405762, -6.42...</td>\n",
       "      <td>testosterone</td>\n",
       "      <td>1953</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>physostigmine</td>\n",
       "      <td>CNC(=O)Oc1ccc2c(c1)[C@]1(C)CCN(C)[C@@H]1N2C</td>\n",
       "      <td>[35.032623291015625, 0.7125436663627625, -1.19...</td>\n",
       "      <td>[-5.117040634155273, -4.733522415161133, -8.03...</td>\n",
       "      <td>physostigmine</td>\n",
       "      <td>1875</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        ATC_LVL5                                             smiles  \\\n",
       "0     novobiocin  CO[C@@H]1[C@@H](OC(N)=O)[C@@H](O)[C@H](Oc2ccc3...   \n",
       "1   vildagliptin      N#C[C@@H]1CCCN1C(=O)CNC12CC3CC(CC(O)(C3)C1)C2   \n",
       "2     benazepril  CCOC(=O)[C@H](CCc1ccccc1)N[C@H]1CCc2ccccc2N(CC...   \n",
       "3   testosterone  CCC(=O)O[C@H]1CC[C@H]2[C@@H]3CCC4=CC(=O)CC[C@]...   \n",
       "4  physostigmine        CNC(=O)Oc1ccc2c(c1)[C@]1(C)CCN(C)[C@@H]1N2C   \n",
       "\n",
       "                                                  mu  \\\n",
       "0  [13.137248039245605, 1.3716269731521606, 0.025...   \n",
       "1  [102.30171203613281, 2.743731737136841, -1.416...   \n",
       "2  [146.45880126953125, 1.0793601274490356, 2.399...   \n",
       "3  [121.25065612792969, 1.7730571031570435, 0.622...   \n",
       "4  [35.032623291015625, 0.7125436663627625, -1.19...   \n",
       "\n",
       "                                                logv      drug_name  Approval  \n",
       "0  [-6.649742126464844, -5.984184741973877, -9.82...     novobiocin      1964  \n",
       "1  [-4.914798259735107, -4.87358283996582, -7.162...   vildagliptin      2007  \n",
       "2  [-4.702251434326172, -4.654637813568115, -6.57...     benazepril      1991  \n",
       "3  [-4.358664035797119, -4.222125053405762, -6.42...   testosterone      1953  \n",
       "4  [-5.117040634155273, -4.733522415161133, -8.03...  physostigmine      1875  "
      ]
     },
     "execution_count": 636,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_smiles.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 637,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1368, 6)"
      ]
     },
     "execution_count": 637,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_smiles.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Morgan Fingerprints (Count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 638,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3879930c2b304cb193d4d2511f8b4f98",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1368), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "fps_c = []\n",
    "for i in tnrange(len(df_smiles)):\n",
    "    smi = df_smiles['smiles'].iloc[i]\n",
    "    mol = Chem.MolFromSmiles(smi)\n",
    "    #fp = dc.feat.rdkit_grid_featurizer.compute_ecfp_features(mol)\n",
    "    fp = AllChem.GetMorganFingerprint(mol,2)\n",
    "    fps_c.append(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 639,
   "metadata": {},
   "outputs": [],
   "source": [
    "# append fingerprints\n",
    "df_smiles['fp_c'] = fps_c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Morgan Fingerprints (Bit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 640,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39359767c4c5435cb6f8f0530701ba2f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1368), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "fps_b = []\n",
    "for i in tnrange(len(df_smiles)):\n",
    "    smi = df_smiles['smiles'].iloc[i]\n",
    "    mol = Chem.MolFromSmiles(smi)\n",
    "    #fp = dc.feat.rdkit_grid_featurizer.compute_ecfp_features(mol)\n",
    "    fp = AllChem.GetMorganFingerprintAsBitVect(mol,2)\n",
    "    fps_b.append(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 641,
   "metadata": {},
   "outputs": [],
   "source": [
    "# append fingerprints\n",
    "df_smiles['fp_b'] = fps_b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Euclidean Drug Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 642,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_dir = './experiments/KDD/kdd_009'\n",
    "checkpoint = 'checkpoint_epoch110.model'\n",
    "config_path = os.path.join(exp_dir, 'configs.json')\n",
    "checkpoint_path = os.path.join(exp_dir, checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 643,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(config_path, 'r') as fp:\n",
    "    configs = json.load(fp)\n",
    "fp.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 644,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'data_dir': './data/fda_drugs',\n",
       " 'data_file': 'smiles_set_clean.smi',\n",
       " 'fda_file': 'all_drugs.smi',\n",
       " 'vocab_file': 'char_set_clean.pkl',\n",
       " 'atc_sim_file': 'drugs_sp_all.csv',\n",
       " 'checkpoint_dir': './experiments/KDD',\n",
       " 'experiment_name': 'kdd_009',\n",
       " 'task': 'vae + atc',\n",
       " 'limit': 0,\n",
       " 'batch_size': 128,\n",
       " 'epochs': 100,\n",
       " 'max_sequence_length': 120,\n",
       " 'learning_rate': 0.0003,\n",
       " 'max_norm': 1000000000000.0,\n",
       " 'wd': 0.0,\n",
       " 'manifold_type': 'Euclidean',\n",
       " 'prior_type': 'Standard',\n",
       " 'num_centroids': 0,\n",
       " 'bidirectional': False,\n",
       " 'num_layers': 1,\n",
       " 'hidden_size': 512,\n",
       " 'latent_size': 64,\n",
       " 'word_dropout_rate': 0.2,\n",
       " 'anneal_function': 'logistic',\n",
       " 'k': 0.51,\n",
       " 'x0': 29.0,\n",
       " 'C': 1.0,\n",
       " 'num_workers': 4,\n",
       " 'logging_steps': 1,\n",
       " 'save_per_epochs': 10,\n",
       " 'new_training': False,\n",
       " 'new_annealing': False,\n",
       " 'checkpoint': 'checkpoint_epoch110.model',\n",
       " 'trained_epochs': 110,\n",
       " 'alpha': 0.0,\n",
       " 'beta': 0.015625,\n",
       " 'gamma': 0.0,\n",
       " 'delta': 11.0,\n",
       " 'nneg': 11,\n",
       " 'fda_prop': 0.2}"
      ]
     },
     "execution_count": 644,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "configs['checkpoint'] = checkpoint\n",
    "configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 645,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EVAE(\n",
      "  (encoder_rnn): GRU(49, 512, batch_first=True)\n",
      "  (decoder_rnn): GRU(49, 512, batch_first=True)\n",
      "  (hidden2mean): Linear(in_features=512, out_features=64, bias=True)\n",
      "  (hidden2logv): Linear(in_features=512, out_features=64, bias=True)\n",
      "  (latent2hidden): Linear(in_features=64, out_features=512, bias=True)\n",
      "  (outputs2vocab): Linear(in_features=512, out_features=49, bias=True)\n",
      "  (RECON): NLLLoss()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = load_model(configs)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 646,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a2aed05677447f9bf388262a3f9fd9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1368), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "mu_e_lst = []\n",
    "for i in tnrange(len(df_smiles)):\n",
    "    smi = df_smiles['smiles'].iloc[i]\n",
    "    mu_e, _= smiles2mean(configs, smi, model)\n",
    "    mu_e_lst.append(mu_e.cpu().detach().numpy().squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 647,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_smiles['mu_e'] = mu_e_lst"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lorentz Drug Embeddings (VAE only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 648,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_dir = './experiments/EXP_TASK/exp_task_010'\n",
    "checkpoint = 'checkpoint_epoch110.model'\n",
    "config_path = os.path.join(exp_dir, 'configs.json')\n",
    "checkpoint_path = os.path.join(exp_dir, checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 649,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(config_path, 'r') as fp:\n",
    "    configs = json.load(fp)\n",
    "fp.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 650,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'data_dir': './data/fda_drugs',\n",
       " 'data_file': 'smiles_set_clean.smi',\n",
       " 'fda_file': 'all_drugs.smi',\n",
       " 'vocab_file': 'char_set_clean.pkl',\n",
       " 'atc_sim_file': 'drugs_sp_all.csv',\n",
       " 'checkpoint_dir': './experiments/EXP_TASK',\n",
       " 'experiment_name': 'exp_task_010',\n",
       " 'task': 'vae',\n",
       " 'limit': 0,\n",
       " 'batch_size': 128,\n",
       " 'epochs': 200,\n",
       " 'max_sequence_length': 120,\n",
       " 'learning_rate': 0.0003,\n",
       " 'max_norm': 1000000000000.0,\n",
       " 'wd': 0.0,\n",
       " 'manifold_type': 'Lorentz',\n",
       " 'prior_type': 'Standard',\n",
       " 'num_centroids': 0,\n",
       " 'bidirectional': False,\n",
       " 'num_layers': 1,\n",
       " 'hidden_size': 512,\n",
       " 'latent_size': 64,\n",
       " 'word_dropout_rate': 0.2,\n",
       " 'anneal_function': 'logistic',\n",
       " 'k': 0.51,\n",
       " 'x0': 29.0,\n",
       " 'C': 1.0,\n",
       " 'num_workers': 4,\n",
       " 'logging_steps': 1,\n",
       " 'save_per_epochs': 5,\n",
       " 'new_training': True,\n",
       " 'new_annealing': True,\n",
       " 'checkpoint': 'checkpoint_epoch110.model',\n",
       " 'trained_epochs': 0,\n",
       " 'alpha': 0.0,\n",
       " 'beta': 0.015625,\n",
       " 'gamma': 0.0,\n",
       " 'delta': 11.0,\n",
       " 'nneg': 11,\n",
       " 'fda_prop': 0.2}"
      ]
     },
     "execution_count": 650,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "configs['checkpoint'] = checkpoint\n",
    "configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 651,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HVAE(\n",
      "  (encoder_rnn): GRU(49, 512, batch_first=True)\n",
      "  (decoder_rnn): GRU(49, 512, batch_first=True)\n",
      "  (hidden2mean): Linear(in_features=512, out_features=64, bias=True)\n",
      "  (hidden2logv): Linear(in_features=512, out_features=64, bias=True)\n",
      "  (latent2hidden): Linear(in_features=65, out_features=512, bias=True)\n",
      "  (outputs2vocab): Linear(in_features=512, out_features=49, bias=True)\n",
      "  (RECON): NLLLoss()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = load_model(configs)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 652,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66d21c8834ed44919c883c156f424ff8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1368), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "mu_vae_lst = []\n",
    "for i in tnrange(len(df_smiles)):\n",
    "    smi = df_smiles['smiles'].iloc[i]\n",
    "    mu_vae, _= smiles2mean(configs, smi, model)\n",
    "    mu_vae_lst.append(mu_vae.cpu().detach().numpy().squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 653,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_smiles['mu_vae'] = mu_vae_lst"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RDKit Descriptors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 654,
   "metadata": {},
   "outputs": [],
   "source": [
    "rdkit_desc = dc.feat.RDKitDescriptors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 655,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17ca3b26fd5a419d943897fb942f1ea9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1368), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "rdkit_desc_lst = []\n",
    "for i in tnrange(len(df_smiles)):\n",
    "    smi = df_smiles['smiles'].iloc[i]\n",
    "    mol = Chem.MolFromSmiles(smi)\n",
    "    features = rdkit_desc._featurize(mol)\n",
    "    rdkit_desc_lst.append(np.array(features))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 656,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_smiles['rdkit_desc'] = rdkit_desc_lst"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation on RepoDB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load RepoDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 657,
   "metadata": {},
   "outputs": [],
   "source": [
    "repoDB = pd.read_csv('./data/drug_repositioning/repoDB.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 658,
   "metadata": {},
   "outputs": [],
   "source": [
    "def str_lower(s):\n",
    "    return s.lower()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 659,
   "metadata": {},
   "outputs": [],
   "source": [
    "repoDB['ATC_LVL5'] = repoDB['drug_name'].map(str_lower)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 660,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>drug_name</th>\n",
       "      <th>drug_id</th>\n",
       "      <th>ind_name</th>\n",
       "      <th>ind_id</th>\n",
       "      <th>NCT</th>\n",
       "      <th>status</th>\n",
       "      <th>phase</th>\n",
       "      <th>DetailedStatus</th>\n",
       "      <th>ATC_LVL5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Lepirudin</td>\n",
       "      <td>DB00001</td>\n",
       "      <td>Heparin-induced thrombocytopenia with thrombosis</td>\n",
       "      <td>C0272275</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Approved</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>lepirudin</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>DB00002</td>\n",
       "      <td>Squamous cell carcinoma of mouth</td>\n",
       "      <td>C0585362</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Approved</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>cetuximab</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>DB00002</td>\n",
       "      <td>Squamous cell carcinoma of nose</td>\n",
       "      <td>C3163899</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Approved</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>cetuximab</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>DB00002</td>\n",
       "      <td>Squamous cell carcinoma of pharynx</td>\n",
       "      <td>C1319317</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Approved</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>cetuximab</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Cetuximab</td>\n",
       "      <td>DB00002</td>\n",
       "      <td>Laryngeal Squamous Cell Carcinoma</td>\n",
       "      <td>C0280324</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Approved</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>cetuximab</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   drug_name  drug_id                                          ind_name  \\\n",
       "0  Lepirudin  DB00001  Heparin-induced thrombocytopenia with thrombosis   \n",
       "1  Cetuximab  DB00002                  Squamous cell carcinoma of mouth   \n",
       "2  Cetuximab  DB00002                   Squamous cell carcinoma of nose   \n",
       "3  Cetuximab  DB00002                Squamous cell carcinoma of pharynx   \n",
       "4  Cetuximab  DB00002                 Laryngeal Squamous Cell Carcinoma   \n",
       "\n",
       "     ind_id  NCT    status phase DetailedStatus   ATC_LVL5  \n",
       "0  C0272275  NaN  Approved   NaN            NaN  lepirudin  \n",
       "1  C0585362  NaN  Approved   NaN            NaN  cetuximab  \n",
       "2  C3163899  NaN  Approved   NaN            NaN  cetuximab  \n",
       "3  C1319317  NaN  Approved   NaN            NaN  cetuximab  \n",
       "4  C0280324  NaN  Approved   NaN            NaN  cetuximab  "
      ]
     },
     "execution_count": 660,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repoDB.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 661,
   "metadata": {},
   "outputs": [],
   "source": [
    "repoDB_merge = repoDB.merge(df_smiles, how='inner', on='ATC_LVL5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 662,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "985"
      ]
     },
     "execution_count": 662,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(repoDB_merge['ATC_LVL5'].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 663,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>drug_name_x</th>\n",
       "      <th>drug_id</th>\n",
       "      <th>ind_name</th>\n",
       "      <th>ind_id</th>\n",
       "      <th>NCT</th>\n",
       "      <th>phase</th>\n",
       "      <th>DetailedStatus</th>\n",
       "      <th>ATC_LVL5</th>\n",
       "      <th>smiles</th>\n",
       "      <th>mu</th>\n",
       "      <th>logv</th>\n",
       "      <th>drug_name_y</th>\n",
       "      <th>Approval</th>\n",
       "      <th>fp_c</th>\n",
       "      <th>fp_b</th>\n",
       "      <th>mu_e</th>\n",
       "      <th>mu_vae</th>\n",
       "      <th>rdkit_desc</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>status</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Approved</th>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "      <td>4738</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Suspended</th>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>257</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "      <td>268</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Terminated</th>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1387</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "      <td>1859</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Withdrawn</th>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>351</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "      <td>449</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            drug_name_x  drug_id  ind_name  ind_id   NCT  phase  \\\n",
       "status                                                            \n",
       "Approved           4738     4738      4738    4738     0      0   \n",
       "Suspended           268      268       268     268   268    268   \n",
       "Terminated         1859     1859      1859    1859  1859   1859   \n",
       "Withdrawn           449      449       449     449   449    449   \n",
       "\n",
       "            DetailedStatus  ATC_LVL5  smiles    mu  logv  drug_name_y  \\\n",
       "status                                                                  \n",
       "Approved                 0      4738    4738  4738  4738         4738   \n",
       "Suspended              257       268     268   268   268          268   \n",
       "Terminated            1387      1859    1859  1859  1859         1859   \n",
       "Withdrawn              351       449     449   449   449          449   \n",
       "\n",
       "            Approval  fp_c  fp_b  mu_e  mu_vae  rdkit_desc  \n",
       "status                                                      \n",
       "Approved        4738  4738  4738  4738    4738        4738  \n",
       "Suspended        268   268   268   268     268         268  \n",
       "Terminated      1859  1859  1859  1859    1859        1859  \n",
       "Withdrawn        449   449   449   449     449         449  "
      ]
     },
     "execution_count": 663,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repoDB_merge.groupby('status').count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split RepoDB dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 664,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use 2000 as the cutoff year\n",
    "cutoff_ind = repoDB_merge['Approval'] < 2000\n",
    "repoDB_merge_train = repoDB_merge[cutoff_ind]\n",
    "repoDB_merge_test = repoDB_merge[~cutoff_ind]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 665,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6201, 19)"
      ]
     },
     "execution_count": 665,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repoDB_merge_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 666,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "787"
      ]
     },
     "execution_count": 666,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(repoDB_merge_train['ATC_LVL5'].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 667,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>drug_name_x</th>\n",
       "      <th>drug_id</th>\n",
       "      <th>ind_name</th>\n",
       "      <th>ind_id</th>\n",
       "      <th>NCT</th>\n",
       "      <th>phase</th>\n",
       "      <th>DetailedStatus</th>\n",
       "      <th>ATC_LVL5</th>\n",
       "      <th>smiles</th>\n",
       "      <th>mu</th>\n",
       "      <th>logv</th>\n",
       "      <th>drug_name_y</th>\n",
       "      <th>Approval</th>\n",
       "      <th>fp_c</th>\n",
       "      <th>fp_b</th>\n",
       "      <th>mu_e</th>\n",
       "      <th>mu_vae</th>\n",
       "      <th>rdkit_desc</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>status</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Approved</th>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "      <td>4320</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Suspended</th>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>210</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "      <td>220</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Terminated</th>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1036</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "      <td>1287</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Withdrawn</th>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>288</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "      <td>374</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            drug_name_x  drug_id  ind_name  ind_id   NCT  phase  \\\n",
       "status                                                            \n",
       "Approved           4320     4320      4320    4320     0      0   \n",
       "Suspended           220      220       220     220   220    220   \n",
       "Terminated         1287     1287      1287    1287  1287   1287   \n",
       "Withdrawn           374      374       374     374   374    374   \n",
       "\n",
       "            DetailedStatus  ATC_LVL5  smiles    mu  logv  drug_name_y  \\\n",
       "status                                                                  \n",
       "Approved                 0      4320    4320  4320  4320         4320   \n",
       "Suspended              210       220     220   220   220          220   \n",
       "Terminated            1036      1287    1287  1287  1287         1287   \n",
       "Withdrawn              288       374     374   374   374          374   \n",
       "\n",
       "            Approval  fp_c  fp_b  mu_e  mu_vae  rdkit_desc  \n",
       "status                                                      \n",
       "Approved        4320  4320  4320  4320    4320        4320  \n",
       "Suspended        220   220   220   220     220         220  \n",
       "Terminated      1287  1287  1287  1287    1287        1287  \n",
       "Withdrawn        374   374   374   374     374         374  "
      ]
     },
     "execution_count": 667,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repoDB_merge_train.groupby('status').count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 668,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1113, 19)"
      ]
     },
     "execution_count": 668,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repoDB_merge_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 669,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "198"
      ]
     },
     "execution_count": 669,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(repoDB_merge_test['ATC_LVL5'].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 670,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>drug_name_x</th>\n",
       "      <th>drug_id</th>\n",
       "      <th>ind_name</th>\n",
       "      <th>ind_id</th>\n",
       "      <th>NCT</th>\n",
       "      <th>phase</th>\n",
       "      <th>DetailedStatus</th>\n",
       "      <th>ATC_LVL5</th>\n",
       "      <th>smiles</th>\n",
       "      <th>mu</th>\n",
       "      <th>logv</th>\n",
       "      <th>drug_name_y</th>\n",
       "      <th>Approval</th>\n",
       "      <th>fp_c</th>\n",
       "      <th>fp_b</th>\n",
       "      <th>mu_e</th>\n",
       "      <th>mu_vae</th>\n",
       "      <th>rdkit_desc</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>status</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Approved</th>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "      <td>418</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Suspended</th>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>47</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "      <td>48</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Terminated</th>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>351</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "      <td>572</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Withdrawn</th>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>63</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "      <td>75</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            drug_name_x  drug_id  ind_name  ind_id  NCT  phase  \\\n",
       "status                                                           \n",
       "Approved            418      418       418     418    0      0   \n",
       "Suspended            48       48        48      48   48     48   \n",
       "Terminated          572      572       572     572  572    572   \n",
       "Withdrawn            75       75        75      75   75     75   \n",
       "\n",
       "            DetailedStatus  ATC_LVL5  smiles   mu  logv  drug_name_y  \\\n",
       "status                                                                 \n",
       "Approved                 0       418     418  418   418          418   \n",
       "Suspended               47        48      48   48    48           48   \n",
       "Terminated             351       572     572  572   572          572   \n",
       "Withdrawn               63        75      75   75    75           75   \n",
       "\n",
       "            Approval  fp_c  fp_b  mu_e  mu_vae  rdkit_desc  \n",
       "status                                                      \n",
       "Approved         418   418   418   418     418         418  \n",
       "Suspended         48    48    48    48      48          48  \n",
       "Terminated       572   572   572   572     572         572  \n",
       "Withdrawn         75    75    75    75      75          75  "
      ]
     },
     "execution_count": 670,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repoDB_merge_test.groupby('status').count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 671,
   "metadata": {},
   "outputs": [],
   "source": [
    "# drug name\n",
    "drug_name_train = list(repoDB_merge_train['ATC_LVL5'].unique())\n",
    "\n",
    "# Lorentz drug embeddings\n",
    "drug_mu_train = []\n",
    "for d in drug_name_train:\n",
    "    df = repoDB_merge_train[repoDB_merge_train['ATC_LVL5'] == d]\n",
    "    drug_mu_train.append(list(df['mu'])[0])\n",
    "\n",
    "# Morgan FP (count) \n",
    "drug_fp_c_train = []\n",
    "for d in drug_name_train:\n",
    "    df = repoDB_merge_train[repoDB_merge_train['ATC_LVL5'] == d]\n",
    "    drug_fp_c_train.append(list(df['fp_c'])[0])\n",
    "\n",
    "# Morgan FP (bit) \n",
    "drug_fp_b_train = []\n",
    "for d in drug_name_train:\n",
    "    df = repoDB_merge_train[repoDB_merge_train['ATC_LVL5'] == d]\n",
    "    drug_fp_b_train.append(list(df['fp_b'])[0])\n",
    "    \n",
    "# Euclidean drug embeddings\n",
    "drug_mu_e_train = []\n",
    "for d in drug_name_train:\n",
    "    df = repoDB_merge_train[repoDB_merge_train['ATC_LVL5'] == d]\n",
    "    drug_mu_e_train.append(list(df['mu_e'])[0])\n",
    "\n",
    "# Lorentz drug embeddings (VAE)\n",
    "drug_mu_vae_train = []\n",
    "for d in drug_name_train:\n",
    "    df = repoDB_merge_train[repoDB_merge_train['ATC_LVL5'] == d]\n",
    "    drug_mu_vae_train.append(list(df['mu_vae'])[0])\n",
    "    \n",
    "# RDKit descriptors\n",
    "drug_rdkit_desc_train = []\n",
    "for d in drug_name_train:\n",
    "    df = repoDB_merge_train[repoDB_merge_train['ATC_LVL5'] == d]\n",
    "    drug_rdkit_desc_train.append(list(df['rdkit_desc'])[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## testing set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 672,
   "metadata": {},
   "outputs": [],
   "source": [
    "# drug name\n",
    "drug_name_test = list(repoDB_merge_test['ATC_LVL5'].unique())\n",
    "\n",
    "# Lorentz drug embeddings\n",
    "drug_mu_test = []\n",
    "for d in drug_name_test:\n",
    "    df = repoDB_merge_test[repoDB_merge_test['ATC_LVL5'] == d]\n",
    "    drug_mu_test.append(list(df['mu'])[0])\n",
    "\n",
    "# Morgan FP (count) \n",
    "drug_fp_c_test = []\n",
    "for d in drug_name_test:\n",
    "    df = repoDB_merge_test[repoDB_merge_test['ATC_LVL5'] == d]\n",
    "    drug_fp_c_test.append(list(df['fp_c'])[0])\n",
    "\n",
    "# Morgan FP (bit) \n",
    "drug_fp_b_test = []\n",
    "for d in drug_name_test:\n",
    "    df = repoDB_merge_test[repoDB_merge_test['ATC_LVL5'] == d]\n",
    "    drug_fp_b_test.append(list(df['fp_b'])[0])\n",
    "    \n",
    "# Euclidean drug embeddings\n",
    "drug_mu_e_test = []\n",
    "for d in drug_name_test:\n",
    "    df = repoDB_merge_test[repoDB_merge_test['ATC_LVL5'] == d]\n",
    "    drug_mu_e_test.append(list(df['mu_e'])[0])\n",
    "\n",
    "# Lorentz drug embeddings (VAE)\n",
    "drug_mu_vae_test = []\n",
    "for d in drug_name_test:\n",
    "    df = repoDB_merge_test[repoDB_merge_test['ATC_LVL5'] == d]\n",
    "    drug_mu_vae_test.append(list(df['mu_vae'])[0])\n",
    "    \n",
    "# RDKit descriptors\n",
    "drug_rdkit_desc_test = []\n",
    "for d in drug_name_test:\n",
    "    df = repoDB_merge_test[repoDB_merge_test['ATC_LVL5'] == d]\n",
    "    drug_rdkit_desc_test.append(list(df['rdkit_desc'])[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pairwise distances between training and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27fda105f3e14cf1a41d383e629d70d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=787), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "drug_pair_dist = [] # Lorentz distance (Chemical + )\n",
    "\n",
    "# Lorentz embedding distance\n",
    "for i in tnrange(len(drug_name_train)):\n",
    "    d_i = drug_name_train[i]\n",
    "    mu_i = drug_mu_train[i]\n",
    "    fp_c_i = drug_fp_c_train[i]\n",
    "    fp_b_i = drug_fp_b_train[i]\n",
    "    mu_e_i = drug_mu_e_train[i]\n",
    "    mu_vae_i = drug_mu_vae_train[i]\n",
    "    rdkit_desc_i = drug_rdkit_desc_train[i]\n",
    "    for j in range(len(drug_name_test)):\n",
    "        d_j = drug_name_test[j]\n",
    "        mu_j = drug_mu_test[j]\n",
    "        fp_c_j = drug_fp_c_test[j]\n",
    "        fp_b_j = drug_fp_b_test[j]\n",
    "        mu_e_j = drug_mu_e_test[j]\n",
    "        mu_vae_j = drug_mu_vae_test[j]\n",
    "        rdkit_desc_j = drug_rdkit_desc_test[j]\n",
    "        \n",
    "        # compute different distances\n",
    "        mu_dist = lor_dist(mu_i, mu_j)\n",
    "        #fp_dist = distance.rogerstanimoto(fp_i,fp_j)\n",
    "        fp_c_dist = 1.0 - DataStructs.TanimotoSimilarity(fp_c_i, fp_c_j)\n",
    "        fp_b_dist = 1.0 - DataStructs.TanimotoSimilarity(fp_b_i, fp_b_j)\n",
    "        mu_e_dist = euc_dist(mu_e_i, mu_e_j)\n",
    "        mu_vae_dist = lor_dist(mu_vae_i, mu_vae_j)\n",
    "        rdkit_desc_dist = angular_dist(rdkit_desc_i, rdkit_desc_j)\n",
    "        drug_pair_dist.append((d_i, d_j, mu_dist, fp_c_dist, fp_b_dist, mu_e_dist, mu_vae_dist, rdkit_desc_dist))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_drug_dist = pd.DataFrame(drug_pair_dist, columns=['drug_train', 'drug_test', 'mu_dist', 'fp_c_dist', \\\n",
    "                                                     'fp_b_dist', 'mu_e_dist', 'mu_vae_dist', 'rdkit_desc_dist'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KNN evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def knn_drug_train(df_drug_dist, drug_test, dist_type, k):\n",
    "    df_sub = df_drug_dist[df_drug_dist['drug_test']==drug_test].sort_values(dist_type)\n",
    "    df_knn = df_sub.head(k)\n",
    "    return list(df_knn['drug_train'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate the kNN outputs dataframe\n",
    "def knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test, dist_type, k):\n",
    "    # loop each drug in test dataset\n",
    "    knn_outputs_all = []\n",
    "    for i in tnrange(len(drug_name_test)):\n",
    "        drug_x = drug_name_test[i]\n",
    "        # select the target drug\n",
    "        df_drug_test = repoDB_merge_test[repoDB_merge_test['ATC_LVL5']==drug_x]\n",
    "        # find KNN in training drug dataset\n",
    "        ind = repoDB_merge_train['ATC_LVL5'].isin(knn_drug_train(df_drug_dist, drug_x, dist_type, k))\n",
    "        df_drug_train = repoDB_merge_train[ind]\n",
    "        # loop each indication of target drug\n",
    "        for index, row in df_drug_test.iterrows():\n",
    "            drug_test = row['ATC_LVL5']\n",
    "            indication_name = row['ind_name']\n",
    "            indication_id = row['ind_id']\n",
    "            status = row['status']\n",
    "            ind_approved = (df_drug_train['ind_id'] == indication_id) & (df_drug_train['status']=='Approved')\n",
    "            ind_terminated = (df_drug_train['ind_id'] == indication_id) & (df_drug_train['status']!='Approved')\n",
    "            cnt_approved = len(df_drug_train[ind_approved]['ATC_LVL5'].unique())\n",
    "            cnt_terminated = len(df_drug_train[ind_terminated]['ATC_LVL5'].unique())\n",
    "            cnt_unknown = k - (cnt_approved + cnt_terminated)\n",
    "            knn_outputs_all.append((drug_test, indication_name, status, cnt_approved, cnt_terminated, cnt_unknown))\n",
    "    df_knn_outputs = pd.DataFrame(knn_outputs_all, columns=['drug_name', 'ind_name', 'status','knn_approved', 'knn_terminated', 'knn_unknown'])\n",
    "    return df_knn_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluation kNN\n",
    "def knn_durg_evaluation(df_knn_outputs, k):\n",
    "    df_DR = df_knn_outputs\n",
    "    # soft votes\n",
    "    label_lst = []\n",
    "    prob_lst = []\n",
    "    for index, row in df_DR.iterrows():\n",
    "        knn_approved = row['knn_approved']\n",
    "        knn_terminated = row['knn_terminated']\n",
    "        knn_unknown = row['knn_unknown']\n",
    "        label = 1.0 if row['status'] == 'Approved' else 0.0\n",
    "        prob = (knn_approved + (0.5 * knn_unknown)) / k\n",
    "        prob_lst.append(prob)\n",
    "        label_lst.append(label)\n",
    "\n",
    "    df_DR['prob'] = prob_lst\n",
    "    df_DR['label'] = label_lst\n",
    "    \n",
    "    # AUROC score\n",
    "    roc = roc_auc_score(np.array(df_DR['label']), np.array(df_DR['prob']))\n",
    "    # AUPRC score\n",
    "    prc = average_precision_score(np.array(df_DR['label']), np.array(df_DR['prob']))\n",
    "    return roc, prc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lorentz Drug Embedding (ATC + ZINC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_dist', 3)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_dist', 5)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 5)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_dist', 7)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 7)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_dist', 9)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 9)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_dist', 11)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 11)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Morgan FP (count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_c_dist', 3)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_c_dist', 5)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 5)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_c_dist', 7)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 7)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_c_dist', 9)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 9)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_c_dist', 11)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 11)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Morgan FP (Bit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_b_dist', 3)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_b_dist', 5)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 5)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_b_dist', 7)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 7)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_b_dist', 9)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 9)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'fp_b_dist', 11)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 11)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Euclidean Drug Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_e_dist', 3)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_e_dist', 5)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 5)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_e_dist', 7)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 7)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_e_dist', 9)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 9)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_e_dist', 11)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 11)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lorentz VAE Only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_vae_dist', 3)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_vae_dist', 5)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 5)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_vae_dist', 7)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 7)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_vae_dist', 9)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 9)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'mu_vae_dist', 11)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 11)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RDKit Descriptor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'rdkit_desc_dist', 3)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'rdkit_desc_dist', 5)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'rdkit_desc_dist', 7)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'rdkit_desc_dist', 9)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knn_outputs = knn_drug_outputs(drug_name_test, repoDB_merge_train, repoDB_merge_test,'rdkit_desc_dist', 11)\n",
    "roc, prc = knn_durg_evaluation(df_knn_outputs, 3)\n",
    "print ('ROC = %0.3f, PRC = %0.3f' %(roc, prc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MolEnv3",
   "language": "python",
   "name": "molenv3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "333px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
